from module_admin.dao.ocr_config_dao import OcrConfigDao
from module_admin.dao.llm_dao import LlmDao
from module_admin.entity.vo.ocr_config_vo import OcrConfigPageQueryModel
import os
from PIL import Image, ImageOps, UnidentifiedImageError
import requests
from config.env import UploadConfig
import time
from openai import OpenAI
import json
import uuid
from config.database import AsyncSessionLocal
from redis.asyncio import Redis
from config.enums import RedisInitKeyConfig
from fastapi import Request
from exceptions.exception import ModelValidatorException
import asyncio
import httpx


class OcrLlmConfig:

    allow_params = [
        "ocr_version",
        "enable_mkldnn",
        "mkldnn_cache_capacity",
        "text_recognition_model_name",
        "text_recognition_model_dir",
        "text_detection_model_name",
        "text_detection_model_dir",
        "doc_orientation_classify_model_name",
        "doc_orientation_classify_model_dir",
        "use_doc_orientation_classify",
        "doc_unwarping_model_name",
        "doc_unwarping_model_dir",
        "use_doc_unwarping",
        "textline_orientation_model_name",
        "textline_orientation_model_dir",
        "use_textline_orientation",
        "use_tensorrt",
        "precision",
    ]

    """
    OCR工具类
    """

    @classmethod
    async def get_ocr_config(cls) -> dict:
        """
        获取OCR配置
        :return: OCR配置对象列表
        """
        async with AsyncSessionLocal() as session:
            ocr_config = await OcrConfigDao.get_OcrConfig_list(
                session, OcrConfigPageQueryModel(), is_page=False
            )
            return cls.parse_ocr_config(ocr_config)

    @classmethod
    def parse_ocr_config(cls, ocr_config: list) -> dict:
        """
        解析OCR配置
        :param ocr_config: OCR配置对象列表
        :return: 解析后的OCR配置字典
        """
        if not ocr_config:
            return {}

        config_dict = {}
        for config in ocr_config:
            if config["confKey"] in cls.allow_params:
                config_dict[config["confKey"]] = config["confValue"]

        return cls.translate_params_type(config_dict)

    @classmethod
    def translate_params_type(cls, ocr_config: dict) -> dict:
        """
        转换OCR配置参数类型
        :param ocr_config: OCR配置字典
        :return: 转换后的OCR配置字典
        """
        if not ocr_config:
            return {}

        translated_config = {}
        for key, value in ocr_config.items():
            if value.lower() in ["true", "false"]:
                translated_config[key] = value.lower() == "true"
            elif value.isdigit():
                translated_config[key] = int(value)
            elif value.lower() == "none":
                translated_config[key] = None
            else:
                translated_config[key] = value

        return translated_config

    @classmethod
    async def init_llm_config(cls, redis: Redis) -> bool:
        """
        初始化LLM的配置
        """
        # 先删除已有的key
        _, keys = await redis.scan(
            0, match="{RedisInitKeyConfig.LLM_CONFIG.key}:*", count=1000
        )
        for key in keys:
            await redis.delete(key)

        async with AsyncSessionLocal() as session:
            llm_config = await LlmDao.get_used_llm_list(session)
            if not llm_config:
                return True
            for config in llm_config:
                await redis.set(
                    f"{RedisInitKeyConfig.LLM_CONFIG.key}:{config['scenario']}",
                    json.dumps(config, ensure_ascii=False, default=str),
                )
            return True

    @classmethod
    async def get_llm_config(cls, scenario: str, redis: Redis) -> dict:
        llm_config_serialized = await redis.get(
            f"{RedisInitKeyConfig.LLM_CONFIG.key}:{scenario}"
        )
        if not llm_config_serialized:
            return None
        return json.loads(llm_config_serialized)

    @classmethod
    async def scenario_check(cls, request: Request):
        scenario = request.query_params.get("scenario", "material_label")
        config = await request.app.state.redis.get(
            f"{RedisInitKeyConfig.LLM_CONFIG.key}:{scenario}"
        )
        if not config:
            raise ModelValidatorException(
                f"scenario -> {scenario}", "unsupported scenario value"
            )


class ImageOptimizer:
    """
    图片优化类，支持JPG/PNG/WEBP格式的压缩优化
    主要功能：
    - 有损压缩（调整质量参数）
    - 自动转换为高效格式
    - 移除EXIF元数据
    - 智能尺寸调整
    - 渐进式渲染（JPEG）
    - 颜色量化（PNG）
    """

    @classmethod
    def image_path(cls, upload: bool = False) -> str:
        """
        生成图片路径
        :param prefix: 文件名前缀
        :return: 生成的图片路径
        """
        date = time.strftime("%Y_%m_%d", time.localtime())
        uuid_str = str(uuid.uuid4()).replace("-", "")
        if upload:
            os.makedirs(f"{UploadConfig.UPLOAD_PATH}/ocr/{date}", exist_ok=True)
            return f"{UploadConfig.UPLOAD_PATH}/ocr/{date}/original_{uuid_str}.jpg"
        else:
            os.makedirs(f"{UploadConfig.DOWNLOAD_PATH}/ocr/{date}", exist_ok=True)
            return f"{UploadConfig.DOWNLOAD_PATH}/ocr/{date}/optimized_{uuid_str}.jpg"

    def __init__(self, input_path):
        """
        初始化优化器
        :param input_path: 输入图片路径
        :param output_path: 输出图片路径（默认覆盖原文件）
        """
        self.input_path = input_path
        self.output_path = UploadConfig.DOWNLOAD_PATH
        self.image = None
        self.original_format = None
        self.ms = 0

    def _open_image(self):
        """打开图片并检测格式"""
        try:
            if self.input_path.startswith("http"):
                response = requests.get(self.input_path)
                response.raise_for_status()
                self.input_path = self.image_path(upload=True)
                with open(self.input_path, "wb") as f:
                    f.write(response.content)
                f.close()

            self.output_path = self.image_path(upload=False)
            self.image = Image.open(self.input_path)
            self.original_format = self.image.format
        except (FileNotFoundError, UnidentifiedImageError) as e:
            raise ValueError(f"无法打开图片: {str(e)}")

    def _should_convert_to_webp(self, target_quality):
        """
        判断是否应转换为WEBP格式
        规则：当原格式为PNG且质量要求<90时转换
        """
        return self.original_format in ["PNG", "APNG"] and target_quality < 90

    def _auto_rotate(self):
        """根据EXIF方向标签自动旋转图片"""
        try:
            self.image = ImageOps.exif_transpose(self.image)
        except Exception:
            pass  # EXIF数据不存在时忽略

    def _resize_large_images(self, max_dimension=2500):
        """调整超大图片尺寸"""
        width, height = self.image.size
        if max(width, height) > max_dimension:
            ratio = max_dimension / max(width, height)
            new_size = (int(width * ratio), int(height * ratio))
            self.image = self.image.resize(new_size, Image.LANCZOS)

    def optimize(
        self, quality=75, remove_metadata=True, progressive=True, max_dimension=None
    ):
        """
        执行图片优化
        :param quality: 压缩质量 (1-100)
        :param remove_metadata: 是否移除元数据
        :param progressive: 是否使用渐进式JPEG
        :param max_dimension: 最大尺寸限制（长边像素）
        :return: 优化后的文件大小（字节）
        """
        start_time = time.time()

        # 打开并预处理图片
        self._open_image()
        self._auto_rotate()

        # 调整超大图片
        if max_dimension:
            self._resize_large_images(max_dimension)

        # 准备保存参数
        save_params = {"quality": quality}

        # 格式转换决策
        output_format = (
            "WEBP" if self._should_convert_to_webp(quality) else self.original_format
        )

        # 设置格式特定参数
        if output_format == "JPEG":
            save_params["progressive"] = progressive
            save_params["optimize"] = True
        elif output_format in ["PNG", "APNG"]:
            save_params["compress_level"] = 9 - int(quality / 15)  # 质量映射到压缩级别

        # 移除元数据
        if remove_metadata:
            save_params["exif"] = b""  # 清空EXIF数据
            self.image.info.pop("icc_profile", None)  # 移除ICC配置文件

        # 保存优化后的图片
        self.image.save(self.output_path, format=output_format, **save_params)
        self.ms = int((time.time() - start_time) * 1000)

        # 返回优化后文件大小
        return os.path.getsize(self.output_path)

    def get_optimization_report(self):
        """获取优化报告"""
        orig_size = os.path.getsize(self.input_path)
        opt_size = os.path.getsize(self.output_path)

        return {
            "original_size": orig_size,
            "optimized_size": opt_size,
            "savings": orig_size - opt_size,
            "reduction_percent": round((1 - opt_size / orig_size) * 100, 1),
            "optimization_time_ms": self.ms,
            "output_path": self.output_path,
        }


class OcrLlmChat:
    @classmethod
    async def chat(cls, user_prompt: str, llm_config: dict) -> str:
        """
        发送聊天请求
        :param user_prompt: 用户输入
        :return: DeepSeek API响应内容
        """
        # todo 查找对应模型的配置
        client = OpenAI(
            api_key=llm_config["apiKey"],
            base_url=llm_config["baseUrl"],
        )

        if not llm_config["systemPrompt"]:
            raise ValueError("LLM系统提示词未配置")

        try:
            response = client.chat.completions.create(
                model=llm_config["modelName"],
                messages=[
                    {"role": "system", "content": llm_config["systemPrompt"]},
                    {"role": "user", "content": user_prompt},
                ],
                stream=False,
                timeout=60,  # SDK 内部请求超时时间
                extra_body={"enable_thinking": False},
            )
            return response.choices[0].message.content.strip()
        except (httpx.TimeoutException, asyncio.TimeoutError) as e:
            # 专门处理超时
            raise TimeoutError(f"调用 LLM 接口超时: {e}") from e
        except Exception as e:
            # 兜底处理其他异常
            raise RuntimeError(f"调用 LLM 接口失败: {e}") from e

    @classmethod
    async def chat_for_label(cls, user_prompt: list[str], llm_config: dict) -> str:
        """
        发送聊天请求，专门用于标签识别
        :param user_prompt: 用户输入
        :return: DeepSeek API响应内容
        """
        up = json.dumps(user_prompt, ensure_ascii=False)
        data = await cls.chat(up, llm_config)
        return data.strip("```json").strip("```").replace("\n", "").replace(" ", "")
